{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the raw data into numpy nd-array\n",
    "\n",
    "data = np.loadtxt('higgs/data.csv', delimiter=',', dtype=np.float32)\n",
    "print('Raw data:', data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the raw data to features, labels and weights\n",
    "\n",
    "features, labels, weights = data[:, 1:-2], data[:, -1], data[:, -2]\n",
    "print('Features', features.shape)\n",
    "print('Labels', labels.shape)\n",
    "print('Weights', weights.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the signal and background distributions for each feature\n",
    "\n",
    "signal_features, signal_weights = features[labels == 1], weights[labels == 1]\n",
    "background_features, background_weights = features[labels == 0], weights[labels == 0]\n",
    "print(signal_features.shape, background_features.shape)\n",
    "\n",
    "plt.figure(figsize=(20, 20))\n",
    "for i in range(30):\n",
    "    plt.subplot(6, 5, i + 1)\n",
    "    plt.hist(signal_features[:, i], bins=50, weights=signal_weights, density=True, color='r', histtype='step', label='s')\n",
    "    plt.hist(background_features[:, i], bins=50, weights=background_weights, density=True, color='b', histtype='step', label='b')\n",
    "    plt.title(str(i))\n",
    "    plt.yticks([])\n",
    "    plt.legend()\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Permutate data\n",
    "\n",
    "indices = np.random.permutation(len(features))\n",
    "print(indices)\n",
    "\n",
    "features, labels, weights = features[indices], labels[indices], weights[indices]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split dataset\n",
    "\n",
    "number_train_examples = int(len(features) * 0.5)\n",
    "print('number of training examples', number_train_examples)\n",
    "print('number of validation examples', len(features) - number_train_examples)\n",
    "\n",
    "train_features, train_labels, train_weights = features[:number_train_examples], labels[:number_train_examples], weights[:number_train_examples]\n",
    "validation_features, validation_labels, validation_weights = features[number_train_examples:], labels[number_train_examples:], weights[number_train_examples:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.tree import DecisionTreeClassifier\n",
    "from sklearn.ensemble import AdaBoostClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build model\n",
    "\n",
    "dt = DecisionTreeClassifier(criterion='gini', max_depth=4)\n",
    "bdt = AdaBoostClassifier(dt, n_estimators=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train\n",
    "\n",
    "bdt.fit(train_features, train_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make predictions on the validation set\n",
    "\n",
    "scores = bdt.predict_proba(validation_features)[:, 1]\n",
    "labels = validation_labels\n",
    "weights = validation_weights\n",
    "\n",
    "print(scores.shape)\n",
    "print(scores)\n",
    "print(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot score distribution\n",
    "\n",
    "signal_scores, signal_weights = scores[labels == 1], weights[labels == 1]\n",
    "background_scores, background_weights = scores[labels == 0], weights[labels == 0]\n",
    "\n",
    "plt.hist(signal_scores, bins=50, weights=signal_weights, density=True, histtype='step', color='r', label='s')\n",
    "plt.hist(background_scores, bins=50, weights=background_weights, density=True, histtype='step', color='b', label='b')\n",
    "plt.xlim(0, 1)\n",
    "plt.ylim(bottom=0)\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "\n",
    "# plt.ylim(0, 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build table [threshold, s, b, Z]\n",
    "\n",
    "# Sort the predictions by score descendingly\n",
    "indices = np.argsort(scores)[::-1]\n",
    "scores, labels, weights = scores[indices], labels[indices], weights[indices]\n",
    "\n",
    "table = np.empty([len(scores) + 1, 4])\n",
    "s, b = 0, 0\n",
    "table[0] = [1, s, b, 0]\n",
    "\n",
    "for i in range(len(scores)):\n",
    "    if labels[i] == 1:\n",
    "        s += weights[i]\n",
    "    else:\n",
    "        b += weights[i]\n",
    "    Z = np.sqrt(2 * ((s + b + 10) * np.log(1 + s / (b + 10)) - s))\n",
    "\n",
    "    table[i + 1] = [scores[i], s, b, Z]\n",
    "\n",
    "np.savetxt('table.txt', table, fmt='%8.3f')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 6))\n",
    "\n",
    "plt.subplot(2, 2, 1)\n",
    "plt.plot(table[:, 0], table[:, 1])\n",
    "plt.xlim(0, 1)\n",
    "plt.ylim(bottom=0)\n",
    "plt.grid()\n",
    "plt.xlabel('Score threshold')\n",
    "plt.ylabel('Number of signal events')\n",
    "\n",
    "plt.subplot(2, 2, 2)\n",
    "plt.plot(table[:, 0], table[:, 2])\n",
    "plt.xlim(0, 1)\n",
    "plt.ylim(bottom=0)\n",
    "plt.grid()\n",
    "plt.xlabel('Score threshold')\n",
    "plt.ylabel('Number of background events')\n",
    "\n",
    "plt.subplot(2, 2, 3)\n",
    "plt.plot(table[:, 2] / table[-1, 2], table[:, 1] / table[-1, 1])\n",
    "plt.xlim(0, 1)\n",
    "plt.ylim(0, 1)\n",
    "plt.grid()\n",
    "plt.xlabel('Background efficiency (false positive rate)')\n",
    "plt.ylabel('Signal efficiency (true positive rate)')\n",
    "\n",
    "plt.subplot(2, 2, 4)\n",
    "plt.plot(table[:, 0], table[:, 3])\n",
    "plt.xlim(0, 1)\n",
    "plt.ylim(bottom=0)\n",
    "plt.grid()\n",
    "plt.xlabel('Score threshold')\n",
    "plt.ylabel('Significance')\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
